'''
Created on Dec 9, 2009

@author: mkiyer
'''
import unittest

import veggie.alignment.sam as sam
import pysam
import numpy as np

def make_read(pos, read_length, is_reverse=False):
    read = type('fakeread', (object,), dict())
    read.pos = pos
    read.rlen = read_length
    read.is_duplicate = False
    read.is_reverse = is_reverse
    #X = type('X', (object,), dict(a=1))
    return read

DTYPE = np.uint32
test_kwargs = {'unique_only': False,
               'merge_strands': True,
               'max_rlen': 2048,
               'chunk_size': 8192,
               'dtype': DTYPE}
    
class Test(unittest.TestCase):

    def test_start(self):        
        # test just one read
        read_positions = [5]
        reads = [make_read(x, 5, False) for x in read_positions]
        expected = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
                            dtype=DTYPE)         
        # pileup the reads chunks at a time        
        arr = sam.pileup_reads(reads, 0, 15, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))        
        # test start not less than end
        raised = False
        self.assertRaises(AssertionError, sam.pileup_reads, reads, 1, 0)
        # no reads in range - reads outside
        read_positions = [15]
        reads = [make_read(x, 5, False) for x in read_positions]
        arr = sam.pileup_reads(reads, 0, 15, **test_kwargs)
        self.assert_(np.array_equal(arr, np.zeros(15, dtype=DTYPE)))
        # chrom start after first read
        read_positions = [10]
        reads = [make_read(x, 5, False) for x in read_positions]
        arr = sam.pileup_reads(reads, 15, 20, **test_kwargs)
        self.assert_(np.array_equal(arr, [0, 0, 0, 0, 0]))
        # chrom start between first and last read
        read_positions = [10, 20]
        reads = [make_read(x, 5, False) for x in read_positions]
        arr = sam.pileup_reads(reads, 15, 20, **test_kwargs)
        self.assert_(np.array_equal(arr, [0, 0, 0, 0, 0]))
        # chrom start halfway inside first/last reads
        read_positions = [10, 15]
        reads = [make_read(x, 5, False) for x in read_positions]
        arr = sam.pileup_reads(reads, 13, 17, **test_kwargs)
        self.assert_(np.array_equal(arr, [1, 1, 1, 1]))
    
    def test_sorted(self):
        read_positions = [5,6,7,8,9,0,1,2,3,4]
        read_length = 10
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
        #arr = sam.pileup_reads(reads, 0, 19, **test_kwargs)
        self.assertRaises(AssertionError, sam.pileup_reads, reads, 0, 19, **test_kwargs)

    def test_consecutive(self):
        read_positions = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        read_length = 10
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
        arr = sam.pileup_reads(reads, 0, 19, **test_kwargs)
        self.assert_(np.array_equal(arr, expected)) 

    def test_repeated(self):
        "Test that the coverage of a stack of duplicate reads is correct"
        read_positions = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
        read_length = 10
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = 10 * np.ones(10, dtype=DTYPE)
        arr = sam.pileup_reads(reads, 0, 10, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))

    def test_overlap_locations(self):
        # 1
        read_positions = [0, 1]
        read_length = 5
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = [1, 2, 2, 2, 2, 1]
        arr = sam.pileup_reads(reads, 0, 6, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        # 2
        read_positions = [0, 2]
        read_length = 5
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = [1, 1, 2, 2, 2, 1, 1]
        arr = sam.pileup_reads(reads, 0, 7, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        # 3
        read_positions = [0, 3]
        read_length = 5
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = [1, 1, 1, 2, 2, 1, 1, 1]
        arr = sam.pileup_reads(reads, 0, 8, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        # 4
        read_positions = [0, 4]
        read_length = 5
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = [1, 1, 1, 1, 2, 1, 1, 1, 1]
        arr = sam.pileup_reads(reads, 0, 9, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        # 5
        read_positions = [0, 5]
        read_length = 5
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
        arr = sam.pileup_reads(reads, 0, 10, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        # 6
        read_positions = [0, 6]
        read_length = 5
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = [1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1]
        arr = sam.pileup_reads(reads, 0, 11, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))

    def test_large_span(self):
        read_positions = [100, 200, 250, 300]
        read_length = 5
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = np.zeros(300, dtype=DTYPE)
        expected[100:105] = 1
        expected[200:205] = 1
        expected[250:255] = 1
        arr = sam.pileup_reads(reads, 0, 300, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))

    def test_chrom_end(self):
        # some arbitrary huge chromosome length
        chrom_length = 247249719
        read_positions = [chrom_length - 1000]
        read_length = 2000
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = np.ones(1000, dtype=DTYPE)
        arr = sam.pileup_reads(reads, chrom_length - 1000, chrom_length, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))

        mykwargs = {'unique_only': False,
                    'merge_strands': True,
                    'max_rlen': 30,
                    'chunk_size': 31,
                    'dtype': DTYPE}
        read_length = 30
        chrom_length = 247249719
        read_position = 247249700
        read = make_read(read_position, read_length, False)     
        expected = np.ones(chrom_length - read_position, dtype=DTYPE)
        arr = sam.pileup_reads([read], read_position, chrom_length, **mykwargs)
        self.assert_(np.array_equal(arr, expected))
        # try with longer fragment length
        arr = sam.pileup_reads([read], read_position, chrom_length, fragment_length=1000, **mykwargs)
        self.assert_(np.array_equal(arr, expected))
        indexes = np.flatnonzero(arr)
        self.assertEqual(read_position + 1 + indexes[-1], chrom_length)
        # try on negative strand
        read = make_read(read_position, read_length, True)     
        arr = sam.pileup_reads([read], read_position, chrom_length, fragment_length=1000, **mykwargs)
        self.assert_(np.array_equal(arr, expected))
        indexes = np.flatnonzero(arr)
        self.assertEqual(read_position + 1 + indexes[-1], chrom_length)
        return

    def test_boundaries(self):
        read_positions = [100, 200]
        read_length = 5
        reads = [make_read(x, read_length, False) for x in read_positions]
        expected = [0, 0, 0]
        arr = sam.pileup_reads(reads, 1000, 1003, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        arr = sam.pileup_reads(reads, 1000, 1003, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        arr = sam.pileup_reads(reads, 5, 8, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))

    def test_fragment_length(self):
        '''
        test pileup of reads on both strands
        '''
        pos_strand_reads = [100]
        neg_strand_reads = [100]
        read_length = 5
        reads = [make_read(x, read_length, False) for x in pos_strand_reads]
        reads.extend([make_read(x, read_length, True) for x in neg_strand_reads])
        expected = [0,0,0,0,0,2,2,2,2,2,0,0,0,0,0]
        arr = sam.pileup_reads(reads, 95, 110, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        # try with different fragment lengths
        fragment_length = 10
        expected = [1,1,1,1,1,2,2,2,2,2,1,1,1,1,1]
        arr = sam.pileup_reads(reads, 95, 110, fragment_length=fragment_length, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        fragment_length = 1
        expected = [0,0,0,0,0,1,0,0,0,1,0,0,0,0,0]
        arr = sam.pileup_reads(reads, 95, 110, fragment_length=fragment_length, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        fragment_length = 2
        expected = [0,0,0,0,0,1,1,0,1,1,0,0,0,0,0]
        arr = sam.pileup_reads(reads, 95, 110, fragment_length=fragment_length, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        fragment_length = 3
        expected = [0,0,0,0,0,1,1,2,1,1,0,0,0,0,0]
        arr = sam.pileup_reads(reads, 95, 110, fragment_length=fragment_length, **test_kwargs)
        self.assert_(np.array_equal(arr, expected))
        return

    def test_chunk_size(self):
        '''
        test pileup when multiple chunks are needed
        '''
        mykwargs = {'unique_only': False,
                    'merge_strands': True,
                    'max_rlen': 10,
                    'dtype': DTYPE}
        read_length = 10
        read_positions = [0, 0, 5, 5, 10, 10, 15, 15, 20, 20, 25, 25, 30, 30] 
        pos_reads = [make_read(x, read_length, False) for x in read_positions]        
        neg_reads = [make_read(x, read_length, True) for x in read_positions]        
        expected = [2,2,2,2,2,4,4,4,4,4,
                    4,4,4,4,4,4,4,4,4,4,
                    4,4,4,4,4,4,4,4,4,4,
                    4,4,4,4,4,2,2,2,2,2]
        # try all different chunk sizes
        for chunk_size in xrange(10, 40):
            arr = sam.pileup_reads(pos_reads, 0, 40, chunk_size=chunk_size, **mykwargs)
            self.assert_(np.array_equal(arr, expected))
            arr = sam.pileup_reads(neg_reads, 0, 40, chunk_size=chunk_size, **mykwargs)
            self.assert_(np.array_equal(arr, expected))
        return


        
if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.main()